/*  电机编号
            头

     9  5       4 8
        1       0
          |dog|
          |dog|
          |dog|
        3       2
     11 7       6 10
*/

#include <iostream>
#include <string>
#include <thread>
#include <atomic>
#include <iomanip>
#include <chrono>

#include <qpOASES.hpp>
#include <Eigen/Dense>

#include "all.h"
#include "func.h"
#include "dog_solve.h"

#include <webots/Robot.hpp>
#include <webots/Motor.hpp>
#include <webots/PositionSensor.hpp>
#include <webots/InertialUnit.hpp>
#include <webots/Gyro.hpp>
#include <webots/Accelerometer.hpp>
#include <webots/Keyboard.hpp>

using namespace std;
using namespace Eigen;
using namespace webots;

// #define NUMBER_OF_LEDS 8
#define NUMBER_OF_JOINTS 12
// #define NUMBER_OF_CAMERAS 5

#define TIME_STEP 8
#define CYCLE_10 10
#define CYCLE_2 2

// *********test flags**************
// #define TEST_THREAD_TRAJECTORY_GENERATE
#define TEST_THREAD_STATE_ESTIMATE

static const char *motor_names[NUMBER_OF_JOINTS] = {
    "front right shoulder abduction motor", // 右前肩外展电机
    "front left shoulder abduction motor",  // 左前肩外展电机
    "rear right shoulder abduction motor",  // 右后肩外展电机
    "rear left shoulder abduction motor",   // 左后肩外展电机
    "front right shoulder rotation motor",  // 右前肩旋转电机
    "front left shoulder rotation motor",   // 左前肩旋转电机
    "rear right shoulder rotation motor",   // 右后肩旋转电机
    "rear left shoulder rotation motor",    // 左后肩旋转电机
    "front right elbow motor",              // 右前肘电机
    "front left elbow motor",               // 左前肘电机
    "rear right elbow motor",               // 右后肘电机
    "rear left elbow motor"                 // 左后肘电机
};

static Motor *motors[NUMBER_OF_JOINTS];
static PositionSensor *ps[NUMBER_OF_JOINTS];

Matrix<double, 13 * PREDICTION_HORIZON, 1> *vecD;
Matrix<double, 13, 1> *x_0;
Matrix<double, 18, 1> *x_k;
Matrix<double, PREDICTION_HORIZON, 1> *psi_k_d;
Matrix<double, 3 * PREDICTION_HORIZON, 4> *r_k_ix;
Matrix<double, 12, 1> *f_MPC;
Matrix<double, 4, 3> *tau_motor;
Matrix<double, 4, 3> *q_motor;

int is_stand[4] = {1, 1, 1, 1};

// Matrix

void init_global_variables() {
  // vecD.store(new Matrix<double, 13 * PREDICTION_HORIZON, 1>());
  // x_0.store(new Matrix<double, 13, 1>());
  // x_k.store(new Matrix<double, 18, 1>());
  // psi_k_d.store(new Matrix<double, PREDICTION_HORIZON, 1>());
  // r_k_ix.store(new Matrix<double, 3 * PREDICTION_HORIZON, 4>());
  vecD = new Matrix<double, 13 * PREDICTION_HORIZON, 1>();
  x_0 = new Matrix<double, 13, 1>();
  x_k = new Matrix<double, 18, 1>();
  psi_k_d = new Matrix<double, PREDICTION_HORIZON, 1>();
  r_k_ix = new Matrix<double, 3 * PREDICTION_HORIZON, 4>();
  f_MPC = new Matrix<double, 12, 1>();
  tau_motor = new Matrix<double, 4, 3>();
  q_motor = new Matrix<double, 4, 3>();
}

void thread_trajectory_generate(Robot *robot, TrajectoryGenerator *tg,
                                StateEstimator *se, GaitPlanner *gp) {

  Keyboard kb;
  kb.enable(16);
  cout << "start thread_trajectory_generate!" << endl;

  // double now = robot->getTime(); // in seconds
  // double prev = now;
  double v_x = 0;
  double v_y = 0;
  double omega_z = 0;

  while (robot->step(4) != -1) {
    int key = kb.getKey();
    // cout<<key<<endl;

    switch (key) {
    case Keyboard::UP:
      cout << "FORWARD" << endl;
      v_x = 0.1;
      v_y = 0;
      omega_z = 0;
      break;
    case Keyboard::DOWN:
      cout << "BACKWARD" << endl;
      v_x = -0.1;
      v_y = 0;
      omega_z = 0;
      break;
    case Keyboard::LEFT:
      cout << "LEFT" << endl;
      v_x = 0;
      v_y = 0;
      omega_z = 0.1;
      break;
    case Keyboard::RIGHT:
      cout << "RIGHT" << endl;
      v_x = 0;
      v_y = 0;
      omega_z = -0.1;
      break;
    default:
      v_x = 0;
      v_y = 0;
      omega_z = 0;
      break;
    }

    tg->setState(se->getp_stend_O(gp->gets_fai()));
    tg->setDesired(v_x, v_y, omega_z);
    tg->calculateAll(se->getR_B_O());
    tg->getD(*vecD);
    tg->getpsi_k_d(*psi_k_d);

    for (int i = 0; i < 4; i++) {
      if (gp->gets_fai()(i) == 1) {
        r_k_ix->col(i) =
            se->getp_foot().row(i).transpose().replicate(1, PREDICTION_HORIZON);
      } else {
        r_k_ix->col(i) = gp->getp_swend_i_O(i).replicate(1, PREDICTION_HORIZON);
      }
    }

    for (int i = 0; i < 4; i++) {
      is_stand[i] = gp->gets_fai()(i);
    }

#ifdef TEST_THREAD_TRAJECTORY_GENERATE
    for (int j = 0; j < PREDICTION_HORIZON; j++) {
      cout << "-----------------------------第 " << j + 1
           << " 周期----------------------------------------" << endl;
      cout << setw(15) << "期望位置:" << std::setw(15)
           << vecD->segment(13 * j + 3, 3).transpose() << endl;
      cout << setw(15) << "期望速度:" << std::setw(15)
           << vecD->segment(13 * j + 9, 3).transpose() << endl;
      cout << setw(15) << "期望角度:" << std::setw(15)
           << vecD->segment(13 * j + 0, 3).transpose() << endl;
      cout << setw(15) << "期望角速度:" << std::setw(15)
           << vecD->segment(13 * j + 6, 3).transpose() << endl;
    }
#endif
  }
}

void thread_state_estimate(Robot *robot, StateEstimator *se, GaitPlanner *gp) {

  cout << "start thread_state_estimate!" << endl;
  // 初始化传感器
  for (int i = 0; i < NUMBER_OF_JOINTS; i++) {
    motors[i] = robot->getMotor(motor_names[i]);
    motors[i]->setPosition(0);
    ps[i] = motors[i]->getPositionSensor();
    ps[i]->enable(TIME_STEP);
  }

  InertialUnit *iu = robot->getInertialUnit("inertial unit");
  iu->enable(TIME_STEP);
  Gyro *gyro = robot->getGyro("gyro");
  gyro->enable(TIME_STEP);
  Accelerometer *acc = robot->getAccelerometer("accelerometer");
  acc->enable(TIME_STEP);

  // 获取IMU数据
  Vector3d Eulerangles = Map<Vector3d>(const_cast<double *>(
      iu->getRollPitchYaw())); // BUG: IMU零系不是通常理解的z轴为偏航
  Vector3d a_original = Map<Vector3d>(const_cast<double *>(acc->getValues()));
  Vector3d omega_original =
      Map<Vector3d>(const_cast<double *>(gyro->getValues()));

  // 编码器获取电机位置和速度
  Matrix<double, 4, 3> q_j_i;
  Matrix<double, 4, 3> q_dot_j_i;

  while (robot->step(4) != -1) {
    for (int i = 0; i < 4; i++) {
      q_j_i(i, 0) = ps[3 * i]->getValue();
      q_j_i(i, 1) = ps[3 * i + 1]->getValue();
      q_j_i(i, 2) = ps[3 * i + 2]->getValue();
      q_dot_j_i(i, 0) = motors[3 * i]->getVelocity();
      q_dot_j_i(i, 1) = motors[3 * i + 1]->getVelocity();
      q_dot_j_i(i, 2) = motors[3 * i + 2]->getVelocity();
    }

    se->setIMUdata(Eulerangles, a_original, omega_original);
    se->setEncoderdata(q_j_i, q_dot_j_i);
    se->calculateAll();
    se->kalman_update();

    *x_k = se->getx_k();

    x_0->segment(0, 3) = Eulerangles;
    x_0->segment(3, 3) = x_k->segment(0, 3);
    x_0->segment(6, 3) = omega_original;
    x_0->segment(9, 3) = x_k->segment(3, 3);
    (*x_0)(12) = -9.8;

#ifdef TEST_THREAD_STATE_ESTIMATE
    // cout<<"x_0: "<<x_0->transpose()<<endl;
    cout << "------------------x_0-----------------:" << endl;
    cout << "位置:" << x_0->segment(3, 3).transpose() << endl;
    cout << "速度:" << x_0->segment(9, 3).transpose() << endl;
    cout << "角度: " << x_0->segment(0, 3).transpose() << endl;
    cout << "角速度: " << x_0->segment(6, 3).transpose() << endl;
#endif

    for (int i = 0; i < 4; i++) {
      // tau_motor->row(i) =
      // (f_MPC->segment(i, 3)).transpose() * JacobianMatrix(q_j_i.row(i), i);
    }
  }
}

void thread_mpc_solver(Robot *robot, MPCsolver *mpc) {
  robot->step(4);
  cout << "start thread_mpc_solver!" << endl;

  mpc->setD_and_x0(*vecD, *x_0);
  mpc->calculateAll(*psi_k_d, *r_k_ix);
  mpc->qpSolveU_first();

  while (robot->step(4) != -1) {
    mpc->setD_and_x0(*vecD, *x_0);
    mpc->calculateAll(*psi_k_d, *r_k_ix);
    mpc->qpSolveU_next(*f_MPC);
  }
}

void thread_send_command(Robot *robot) {

  cout << "start thread_send_command!" << endl;
  while (robot->step(4) != -1) {
    for (int i = 0; i < 4; i++) {
      if (is_stand[i] == 1) {
        motors[i]->setTorque((*tau_motor)(i, 0));
        motors[i + 4]->setTorque((*tau_motor)(i, 1));
        motors[i + 8]->setTorque((*tau_motor)(i, 2));
      } else if (is_stand[i] == 0) {
        motors[i]->setPosition((*q_motor)(i, 0));
        motors[i + 4]->setPosition((*q_motor)(i, 1));
        motors[i + 8]->setPosition((*q_motor)(i, 2));
      }
    }
  }
}

void trot_balanced() {

  Robot *robot = new Robot();
  GaitPlanner *gp = new GaitPlanner();
  TrajectoryGenerator *tg = new TrajectoryGenerator();
  StateEstimator *se = new StateEstimator();
  MPCsolver *mpc = new MPCsolver();

  init_global_variables();
  // thread thread1(thread_trajectory_generate, robot, tg, se, gp);
  thread thread2(thread_state_estimate, robot, se, gp);
  // thread thread3(thread_mpc_solver, robot, mpc);
  // thread thread4(thread_send_command, robot);

  // thread1.join();
  thread2.join();
  // thread3.join();
  // thread4.join();

  delete gp;
  delete tg;
  delete se;
  delete mpc;
  delete robot;
}

int main() {
  trot_balanced();
  return 0;
}
